"""
shell工具
"""
from typing import Optional, List, Dict, Any
from pydantic import Field
from mcp.server import FastMCP
from .container import get_container
import docker
import subprocess


def add_shell_tools(mcp: FastMCP, _docker: docker.DockerClient):

    @mcp.tool()
    async def execute_shell_command(
            name: str = Field(default="python-sandbox", description="容器名称."),
            container_id: Optional[str] = Field(default=None, description="容器id."),
            shell_command: str = Field(default=..., description="shell命令."),
    ) -> Dict[str, Any]:
        """在容器中执行shell命令"""
        container = get_container(_docker, name, container_id)
        if container is None:
            return {'success': False, 'message': '未找到相应的容器.'}
        command = f"docker exec -i {container.short_id} {shell_command}"
        result = subprocess.run(command, shell=True, capture_output=True)
        if result.returncode == 0:
            return {'success': True, 'message': result.stdout}
        else:
            return {'success': False, 'message': result.stderr.decode()}

    @mcp.tool()
    async def clean_docker_and_images() -> List[Dict[str, Any]]:
        """清理所有状态为Exit和Created的容器，清理所有None镜像. """

        commands = [
            # 停止所有Exit的容器
            """docker ps -a | grep "Exited" | awk '{print $1 }' | xargs docker stop""",
            # 删除所有Exit的容器
            """docker ps -a | grep "Exited" | awk '{print $1 }' | xargs docker rm""",
            # 停止所有状态为Created的容器
            """docker ps -f "status=created" | awk '{print $1 }' | xargs docker stop""",
            # 删除所有状态为Created的容器
            """docker ps -f "status=created" | awk '{print $1 }' | xargs docker rm""",
            # 删除所有none镜像
            """docker images | grep none | awk '{print $3 }' | xargs docker rmi - f"""
        ]
        results = []
        for command in commands:
            res = subprocess.run(command, shell=True, capture_output=True)
            if res.returncode == 0:
                result = {'success': True, 'message': res.stdout}
            else:
                result = {'success': False, 'message': res.stderr.decode()}
            results.append(result)
        return results
